BeginPackage["NVInteraction`",{"MathMR`", "MRData`"}];
$Xdir		= Normalize[{1,-1,0}];
$Ydir		= Normalize[{1,1,-2}];
$Zdir		= Normalize[{1,1,1}];
$QubitPos	= $DiamondR1*$DiamondLatticeConst;
$DiamondZFS	= 2.87(*GHz*);
$DiamondESZFS = 1.4(*GHz*);
$NVn14ZFS = -5.04(*MHz*);
$HFn14vert = 2.3(*MHz*);
$HFn14zz = 2.1(*MHz*);
$HFn15 = -3.05(*MHz*);
$IntConstNN	= $Mu0/4/Pi*$GammaC13*$GammaC13*($Hbar*$Hbar)/$Angstrom/$Angstrom/$Angstrom/($Hbar*1000);(*Unit: k rad / s *)
$IntConstNE	= $Mu0/4/Pi*$GammaC13*$GammaE*($Hbar*$Hbar)/$Angstrom/$Angstrom/$Angstrom/($Hbar*1000);(*Unit: k rad / s *)

NaturalFrame2CrystalFrame	::usage="";
CrystalFrame2NaturalFrame	::usage="";
NaturalFrame2BFrame			::usage="";
BFrame2NaturalFrame			::usage="";
CrystalFrame2BFrame			::usage="";
BFrame2CrystalFrame			::usage="";

QubitHami			::usage="Calculate the NV Center Hamiltonian for given magnetic field.";
ExcitedStateHami	::usage="";
Eigen				::usage="Calculate the Eigensystem of the NV Hamiltonian.";
ExcitedStateEigen	::usage="";
Svalue				::usage="Calculate the mean spin values for a given state and magnetic field.";
MagneticBDirection	::usage="Calculate the direction vector according to the given magnetic field polar and azimuth angle.";
MagBFrame			::usage="Calculate the coordinate frame according to the given magnetic field direction.";
HyperfineField		::usage="Calculate the HF field for given position, state (or mean spin value), and magnetic field.";
HyperfineTensor		::usage="";
MediatedInteraction	::usage="";
DipolarInteraction	::usage="Calculate the dipolar interaction between two spins.";
DipInteractions		::usage="ML"
DipolarInteractionTable	::usage="";
(*=================================================================*)
(*Qubit States*)
Begin["`QubitStates`"];

MagneticBDirection[theta_,phi_]:=
	Module[{res},
		res=Sin[theta]Cos[phi]*$Xdir+Sin[theta]Sin[phi]*$Ydir+Cos[theta]*$Zdir;
		Return[res];
	]

MagBFrame[theta_,phi_]:=
	Module[{res,xdir,ydir,zdir},
		res=If[theta==0.,
				{$Xdir,$Ydir,$Zdir}
				,
				zdir=MagneticBDirection[theta,phi];
				ydir=Cross[$Zdir,zdir]//Normalize;
				xdir=Cross[ydir,zdir]//Normalize;
				{xdir,ydir,zdir}
			];
		Return[res];
	]

QubitHami[b_,theta_,phi_]:=
	Module[{res,omega,omegaX,omegaY,omegaZ,Spin1x,Spin1y,Spin1z},
		{Spin1x,Spin1y,Spin1z}={SpinOperatorX[1],SpinOperatorY[1],SpinOperatorZ[1]};
		
		omega=-b*$GammaE/(2*Pi*10^9);
		omegaX=omega*Sin[theta]Cos[phi];
		omegaY=omega*Sin[theta]Sin[phi];
		omegaZ=omega*Cos[theta];
		
		res=$DiamondZFS*Spin1z.Spin1z+omegaX*Spin1x+omegaY*Spin1y+omegaZ*Spin1z;
		Return[res];
	]

ExcitedStateHami[b_,theta_,phi_,strain_]:=
	Module[{res,omega,omegaX,omegaY,omegaZ,Spin1x,Spin1y,Spin1z},
		{Spin1x,Spin1y,Spin1z}={SpinOperatorX[1],SpinOperatorY[1],SpinOperatorZ[1]};
		
		omega=-b*$GammaE/(2*Pi*10^9);
		omegaX=omega*Sin[theta]Cos[phi];
		omegaY=omega*Sin[theta]Sin[phi];
		omegaZ=omega*Cos[theta];
		
		res=$DiamondESZFS*Spin1z.Spin1z+omegaX*Spin1x+omegaY*Spin1y+omegaZ*Spin1z+strain*(Spin1x.Spin1x-Spin1y.Spin1y);
		Return[res];
	]

Eigen[b_,theta_,phi_]:=
	Module[{res,hami,u,v},
		hami=QubitHami[b,theta,phi];
		{u,v}=Eigensystem[hami];
		v=v[[Ordering[u]]];
		u=Sort[u];
		res={u,v};
		Return[res];
	]

ExcitedStateEigen[b_,theta_,phi_,strain_]:=
	Module[{res,hami,u,v},
		hami=ExcitedStateHami[b,theta,phi,strain];
		{u,v}=Eigensystem[hami];
		v=v[[Ordering[u]]];
		u=Sort[u];
		res={u,v};
		Return[res];
	]

Svalue[state_,b_,theta_,phi_]:=
	Module[{res,u,v,st,svalXYZ,Spin1x,Spin1y,Spin1z},
		{Spin1x,Spin1y,Spin1z}={SpinOperatorX[1],SpinOperatorY[1],SpinOperatorZ[1]};
		{u,v}=Eigen[b,theta,phi];
		st=v[[state]];
		svalXYZ=Chop[{Conjugate[st].(Spin1x.st),Conjugate[st].(Spin1y.st),Conjugate[st].(Spin1z.st)}];
		res=svalXYZ.{$Xdir,$Ydir,$Zdir};
		Return[res];
	]

End[];

(*=================================================================*)
(*Qubit Bath Interactions*)
Begin["`QubitBathInt`"];

NaturalFrame2CrystalFrame[Ncomponents_]:=
	Module[{res, RmatCN},
		RmatCN={$Xdir, $Ydir, $Zdir};
		res=Ncomponents.Transpose[RmatCN];
		Return[res];
	]

CrystalFrame2NaturalFrame[Ccomponents_]:=
	Module[{res,RmatCN},
		RmatCN={$Xdir, $Ydir, $Zdir};
		res=Ccomponents.RmatCN;
		Return[res];
	]

NaturalFrame2BFrame[Ncomponents_,theta_,phi_]:=
	Module[{res,RmatBN},
		RmatBN=MagBFrame[theta,phi];
		res=Ncomponents.Transpose[RmatBN];
		Return[res];
	]

BFrame2NaturalFrame[Bcomponents_,theta_,phi_]:=
	Module[{res,RmatBN},
		RmatBN=MagBFrame[theta,phi];
		res=Bcomponents.RmatBN;
		Return[res];
	]
	
CrystalFrame2BFrame[Ccomponents_,theta_,phi_]:=
	Module[{res,RmatCN,RmatBN,RmatCB},
		RmatCN={$Xdir, $Ydir, $Zdir};
		RmatBN=MagBFrame[theta,phi];
		RmatCB=RmatCN.Transpose[RmatBN];
		res=Ccomponents.RmatCB;
		Return[res];
	]

BFrame2CrystalFrame[Bcomponents_,theta_,phi_]:=
	Module[{res,RmatCN,RmatBN,RmatCB},
		RmatCN={$Xdir, $Ydir, $Zdir};
		RmatBN=MagBFrame[theta,phi];
		RmatCB=RmatCN.Transpose[RmatBN];
		res=Bcomponents.Transpose[RmatCB];
		Return[res];
	]

MediatedInteraction[coord1_,coord2_,state_,b_,theta_,phi_]:=
	Module[{res, Amat1, Amat2,eigVal,eigVec,perturbState,
		a,b1,b2,delta1,delta2,Sab1, Sab2, Smat1,Smat2, Dmat1,Dmat2},
		
		Amat1=HyperfineTensor[coord1,b,theta,phi];
		Amat2=HyperfineTensor[coord2,b,theta,phi];
		
		perturbState=Complement[{1,2,3},{state}];
		{eigVal,eigVec}=Eigen[b,theta,phi];
		a=eigVec[[state]];
		b1=eigVec[[ perturbState[[1]] ]];
		b2=eigVec[[ perturbState[[2]] ]];
		delta1=(eigVal[[state]]-eigVal[[ perturbState[[1]] ]])*2.0*Pi*10^6;
		delta2=(eigVal[[state]]-eigVal[[ perturbState[[2]] ]])*2.0*Pi*10^6;
		
		
		Sab1=CrystalFrame2BFrame[{Conjugate[a].SpinOperatorX[1].b1, Conjugate[a].SpinOperatorY[1].b1,	Conjugate[a].SpinOperatorZ[1].b1},theta,phi];
		Sab2=CrystalFrame2BFrame[{Conjugate[a].SpinOperatorX[1].b2, Conjugate[a].SpinOperatorY[1].b2,	Conjugate[a].SpinOperatorZ[1].b2},theta,phi];
		
		Smat1=Table[Sab1[[i]]*Conjugate[Sab1[[j]]],{i,1,3},{j,1,3}];
		Smat2=Table[Sab2[[i]]*Conjugate[Sab2[[j]]],{i,1,3},{j,1,3}];		
		
		Dmat1=Amat1.Smat1.Amat2;
		Dmat2=Amat1.Smat2.Amat2;
		If[Abs[delta1]>0.0, 
			res = Dmat1/delta1;
			(*Print[delta1//Chop,"\t",Norm[Amat1.Smat1.Amat2/delta1,2]];*)
			,
			res = {{0.,0.,0.},{0.,0.,0.},{0.,0.,0.}}; 
		];
		
		If[Abs[delta2]>0.0,
			res = res + Dmat2/delta2;
			(*Print[delta2//Chop,"\t",Norm[Amat1.Smat1.Amat2/delta2,2]];*)
		];
		
		res = If[coord1!=coord2,
			 2.0*Re[res],
			 {res[[2,3]]-res[[3,2]], res[[3,1]]-res[[1,3]], res[[1,2]]-res[[2,1]]}*I/2.0
		];
		
		Return[res];
	]

HyperfineTensor[coord_,b_,theta_,phi_]:=
	Module[{res, xdir, ydir, zdir, rvec, r, n, nx, ny, nz, mat},
		{xdir,ydir,zdir}=MagBFrame[theta,phi];
		rvec=coord-$QubitPos;
		r=Norm[rvec];
		n=rvec/r;
		
		nx=n.xdir;
		ny=n.ydir;
		nz=n.zdir;
		
		mat={
			{1.0-3.0 nx*nx, -3.0*nx*ny, 	-3.0*nx*nz},
			{-3.0*ny*nx, 	1.0-3.0*ny*ny, 	-3.0*ny*nz},
			{-3.0*nz*nx, 	-3.0*nz*ny, 	1.0-3.0*nz*nz}
			};
		
		res=$IntConstNE/r^3*mat;
		Return[res];
	]

HyperfineField[coord_,state_,b_,theta_,phi_,component_:{1,2,3}]:=
	Module[{res,xdir,ydir,zdir,sval,sx,sy,sz,nx,ny,nz,rvec,r,n,sn,resX,resY,resZ},
		{xdir,ydir,zdir}=MagBFrame[theta,phi];
		sval=Svalue[state,b,theta,phi];
		
		rvec=coord-$QubitPos;
		r=Norm[rvec];
		n=rvec/r;
		
		sn=sval.n;
		sx=sval.xdir;nx=n.xdir;
		sy=sval.ydir;ny=n.ydir;
		sz=sval.zdir;nz=n.zdir;
		
		resX=sx-3.sn*nx;(*k rad / s*)
		resY=sy-3.sn*ny;(*k rad / s*)
		resZ=sz-3.sn*nz;(*k rad / s*)
		
		res=($IntConstNE/r^3*{resX,resY,resZ})[[component]];
		Return[res];
	]/;NumberQ[state]

HyperfineField[coord_,sval_,b_,theta_,phi_,component_:{1,2,3}]:=
	Module[{res,xdir,ydir,zdir,sx,sy,sz,nx,ny,nz,rvec,r,n,sn,resX,resY,resZ},
		{xdir,ydir,zdir}=MagBFrame[theta,phi];
		
		rvec=coord-$QubitPos;
		r=Norm[rvec];
		n=rvec/r;
		
		sn=sval.n;
		sx=sval.xdir;nx=n.xdir;
		sy=sval.ydir;ny=n.ydir;
		sz=sval.zdir;nz=n.zdir;
		
		resX=sx-3.sn*nx;(*k rad / s*)
		resY=sy-3.sn*ny;(*k rad / s*)
		resZ=sz-3.sn*nz;(*k rad / s*)
		
		res=($IntConstNE/r^3*{resX,resY,resZ})[[component]];
		Return[res];
	]/;(Head[sval]==List)&&(Length[sval]==3)

End[];

(*=================================================================*)
(*Intra-Bath Interactions*)
Begin["`IntraBathInt`"];

(*$LINKDipInt=Install[$MRMathLinkPath<>"/MLDipInteractions/MLDipInteractions"];*)

DipolarInteractionTable[coordList_,b_,theta_,phi_,gammaList_,component_:{1,2,3,4,5,6}]:=
	Module[{res,len,lk},
		lk=Install[$MRMathLinkPath<>"/MLDipInteractions/MLDipInteractions"];
		len=Length[coordList];
		res=Partition[#[[component]]& /@ Partition[DipInteractions[coordList//Flatten, gammaList,{b, theta, phi}],6],len];
		Uninstall[lk];
			
		Return[res];
	]

DipolarInteraction[coord1_,coord2_,b_,theta_,phi_,component_:{1,2,3,4,5,6}]:=
	Module[{res,xdir,ydir,zdir,r12vec,r12,n,nx,ny,nz,resXX,resYY,resZZ,
			resXY,resYZ,resZX},
			
		{xdir,ydir,zdir}=MagBFrame[theta,phi];
		r12vec=coord2-coord1;
		r12=Norm[r12vec];
		n=r12vec/r12;
		
		nx=n.xdir;
		ny=n.ydir;
		nz=n.zdir;
		
		resXX=1.0-3.0nx*nx;
		resYY=1.0-3.0ny*ny;
		resZZ=1.0-3.0nz*nz;
		resXY=-3.0nx*ny;
		resYZ=-3.0ny*nz;
		resZX=-3.0nz*nx;
		
		res=($IntConstNN/r12^3*{resXX,resYY,resZZ,resXY,resYZ,resZX})[[component]];
		Return[res];
	]

End[];

(*=================================================================*)
EndPackage[];
